package com.mavenq.fly.websocket.interceptor;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

@Slf4j
public class AuthHandshakeInterceptor implements HandshakeInterceptor {

    private static final String TOKEN_PARAM = "token";
    private static final String USER_ID_PARAM = "userId";

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        // 仅处理 WebSocket 握手
        if (!(request.getHeaders().getUpgrade() != null && "websocket".equalsIgnoreCase(
                request.getHeaders().getUpgrade().toString()))) {
            return true;
        }

        String token = getToken(request);
        String userId = getUserId(request);

        if (StringUtils.isBlank(token)) {
            response.setStatusCode(HttpStatus.UNAUTHORIZED);
            log.warn("Handshake denied: missing token.");
            return false;
        }
        if (StringUtils.isBlank(userId)) {
            response.setStatusCode(HttpStatus.BAD_REQUEST);
            log.warn("Handshake denied: missing userId.");
            return false;
        }

        // 可在此处接入真实鉴权（如 Redis 校验、JWT 校验等）
        if (!"valid-token".equalsIgnoreCase(token)) {
            response.setStatusCode(HttpStatus.UNAUTHORIZED);
            log.warn("Handshake denied: invalid token. userId={}", userId);
            return false;
        }

        attributes.put("token", token);
        attributes.put("userId", userId);
        log.info("Handshake allowed. userId={}, token=***", userId);
        return true;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception exception) {
    }

    private String getToken(ServerHttpRequest request) {
        return request.getURI().getQuery() == null ? null :
                request.getURI().getQuery().split("&")[0].startsWith("token=") ?
                        request.getURI().getQuery().split("&")[0].substring(6) : null;
    }

    private String getUserId(ServerHttpRequest request) {
        return request.getURI().getQuery() == null ? null :
                request.getURI().getQuery().split("&")[0].startsWith("userId=") ?
                        request.getURI().getQuery().split("&")[0].substring(7) : null;
    }
}